"""
清空数据库脚本
运行此脚本将清空所有抽取记录、实体、关系、词典和文档数据
保留用户账户信息
"""
from app import app, db
from models import ExtractionRecord, KnowledgeEntity, KnowledgeRelation, DictionaryTerm, Document
import os

def clear_knowledge_data():
    """清空知识数据，保留用户信息"""
    with app.app_context():
        try:
            # 删除所有关系数据
            deleted_relations = KnowledgeRelation.query.delete()
            print(f"已删除 {deleted_relations} 条关系数据")
            
            # 删除所有实体数据
            deleted_entities = KnowledgeEntity.query.delete()
            print(f"已删除 {deleted_entities} 条实体数据")
            
            # 删除所有抽取记录
            deleted_extractions = ExtractionRecord.query.delete()
            print(f"已删除 {deleted_extractions} 条抽取记录")
            
            # 删除所有词典数据
            deleted_terms = DictionaryTerm.query.delete()
            print(f"已删除 {deleted_terms} 条词典数据")
            
            # 删除所有文档数据和文件
            documents = Document.query.all()
            deleted_docs = 0
            for doc in documents:
                if os.path.exists(doc.file_path):
                    try:
                        os.remove(doc.file_path)
                    except:
                        pass
                deleted_docs += 1
            Document.query.delete()
            print(f"已删除 {deleted_docs} 个文档")
            
            # 提交更改
            db.session.commit()
            print("\n✅ 数据库清空成功！用户账户已保留。")
            
        except Exception as e:
            db.session.rollback()
            print(f"\n❌ 清空失败：{e}")

def clear_all_data():
    """完全清空数据库，包括用户信息"""
    with app.app_context():
        try:
            # 删除所有表数据
            db.drop_all()
            print("已删除所有表")
            
            # 重新创建表
            db.create_all()
            print("已重新创建表")
            
            # 创建默认管理员账户
            from models import User
            admin = User(
                username='admin',
                is_admin=True
            )
            admin.set_password('admin123')
            db.session.add(admin)
            db.session.commit()
            
            print("\n✅ 数据库已完全重置！")
            print("默认管理员账户：admin / admin123")
            
        except Exception as e:
            db.session.rollback()
            print(f"\n❌ 重置失败：{e}")

if __name__ == '__main__':
    import sys
    
    print("=" * 60)
    print("数据库清空工具")
    print("=" * 60)
    print()
    print("请选择操作：")
    print("1. 清空知识数据（保留用户账户）")
    print("2. 完全重置数据库（包括用户账户）")
    print("0. 取消")
    print()
    
    choice = input("请输入选项 (0/1/2): ").strip()
    
    if choice == '1':
        confirm = input("确认清空知识数据？(yes/no): ").strip().lower()
        if confirm == 'yes':
            clear_knowledge_data()
        else:
            print("已取消")
    elif choice == '2':
        confirm = input("确认完全重置数据库？(yes/no): ").strip().lower()
        if confirm == 'yes':
            clear_all_data()
        else:
            print("已取消")
    else:
        print("已取消")

